import requests
import json
import time
import random

BASE_URL = 'http://127.0.0.1:3000'

def create_collection(collection_name):
    url = f"{BASE_URL}/create_collection"
    data = {"name": collection_name}
    response = requests.post(url, json=data)
    if response.status_code == 200:
        print(f"Collection '{collection_name}' created successfully.")
    else:
        print(f"Failed to create collection '{collection_name}': {response.text}")

def bulk_insert(collection_name, num_vectors, dim, batch_size=1000):
    url = f"{BASE_URL}/insert"
    vectors = [
        {
            "id": f"vector_{i}",
            "embedding": [random.random() for _ in range(dim)],
            "metadata": {"metadata": f"vector_{i}"}
        }
        for i in range(num_vectors)
    ]
    start_time = time.time()
    for i in range(0, num_vectors, batch_size):
        batch = vectors[i:i + batch_size]
        data = {
            "collection_name": collection_name,
            "vectors": batch
        }
        response = requests.post(url, json=data)
        if response.status_code != 200:
            print(f"Failed to insert batch {i // batch_size + 1}: {response.text}")
            return
    end_time = time.time()
    print(f"Inserted {num_vectors} vectors into '{collection_name}' in {end_time - start_time:.2f} seconds.")

def search(collection_name, query_vector, k):
    url = f"{BASE_URL}/search"
    data = {
        "collection_name": collection_name,
        "query": query_vector,
        "k": k
    }
    start_time = time.time()
    response = requests.post(url, json=data)
    end_time = time.time()
    if response.status_code == 200:
        results = response.json()
        print(f"Search completed in {end_time - start_time:.4f} seconds. Results: {results}")
    else:
        print(f"Search failed: {response.text}")

def main():
    collection_name = "performance_test"
    num_vectors = 10000  # Number of vectors to insert
    batch_size = 2000
    dim = 128  # Dimensionality of vectors
    k = 10  # Top-k results for search

    # Step 1: Create collection
    create_collection(collection_name)

    # Step 2: Bulk insert vectors
    bulk_insert(collection_name, num_vectors, dim, batch_size)

    # Step 3: Perform search
    query_vector = [random.random() for _ in range(dim)]
    search(collection_name, query_vector, k)

if __name__ == "__main__":
    main()
